from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from datetime import datetime
import json
import logging
from enum import Enum
from typing import Dict, Optional, List
from pydantic import BaseModel
import asyncio

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="设备管理与AI对话服务")


class MessageType(Enum):
    DEVICE_REGISTER = "device_register"
    STATUS_UPDATE = "status_update"
    DATA_REPORT = "data_report"
    AI_CHAT = "ai_chat"
    HEARTBEAT = "heartbeat"


class DeviceStatus(BaseModel):
    deviceId: str
    versionName: Optional[str] = None
    url: Optional[str] = None
    code: Optional[str] = None
    status: str = "active"
    lastSeen: Optional[datetime] = None


class WebSocketMessage(BaseModel):
    type: MessageType
    data: Dict
    timestamp: Optional[datetime] = None


class ConnectionManager:
    def __init__(self):
        # 设备连接管理
        self.active_device_connections: Dict[str, WebSocket] = {}
        # AI对话会话管理
        self.chat_sessions: Dict[str, List[Dict]] = {}

    async def connect_device(self, device_id: str, websocket: WebSocket):
        """连接设备"""
        await websocket.accept()
        self.active_device_connections[device_id] = websocket
        logger.info(f"设备 {device_id} 连接成功")

    def disconnect_device(self, device_id: str):
        """断开设备连接"""
        if device_id in self.active_device_connections:
            del self.active_device_connections[device_id]
        if device_id in self.chat_sessions:
            del self.chat_sessions[device_id]
        logger.info(f"设备 {device_id} 断开连接")

    async def send_personal_message(self, message: dict, device_id: str):
        """向指定设备发送消息"""
        if device_id in self.active_device_connections:
            try:
                await self.active_device_connections[device_id].send_text(json.dumps(message))
            except Exception as e:
                logger.error(f"向设备 {device_id} 发送消息失败: {str(e)}")

    async def broadcast(self, message: dict):
        """广播消息给所有设备"""
        for device_id, websocket in self.active_device_connections.items():
            try:
                await websocket.send_text(json.dumps(message))
            except Exception as e:
                logger.error(f"广播消息到设备 {device_id} 失败: {str(e)}")


class AIChatService:
    """AI对话服务"""

    def __init__(self):
        self.chat_history_limit = 10  # 保存最近10条对话记录

    async def process_chat_message(self, device_id: str, user_message: str, manager: ConnectionManager) -> str:
        """处理AI对话消息"""
        try:
            # 初始化或获取对话历史
            if device_id not in manager.chat_sessions:
                manager.chat_sessions[device_id] = []

            chat_history = manager.chat_sessions[device_id]

            # 添加用户消息到历史
            chat_history.append({
                "role": "user",
                "content": user_message,
                "timestamp": datetime.now().isoformat()
            })

            # 调用AI服务生成回复（这里使用模拟的AI回复）
            ai_response = await self.generate_ai_response(user_message, chat_history)

            # 添加AI回复到历史
            chat_history.append({
                "role": "assistant",
                "content": ai_response,
                "timestamp": datetime.now().isoformat()
            })

            # 限制历史记录长度
            if len(chat_history) > self.chat_history_limit:
                manager.chat_sessions[device_id] = chat_history[-self.chat_history_limit:]

            return ai_response

        except Exception as e:
            logger.error(f"AI对话处理失败: {str(e)}")
            return "抱歉，我现在无法处理您的请求，请稍后再试。"

    async def generate_ai_response(self, user_message: str, history: List[Dict]) -> str:
        """生成AI回复（这里使用模拟实现，可以替换为真实的AI服务）"""
        # 模拟AI处理时间
        await asyncio.sleep(0.5)

        # 简单的规则回复（实际应用中可替换为 OpenAI、文心一言等AI服务）
        user_message_lower = user_message.lower()

        if any(word in user_message_lower for word in ["你好", "hello", "hi"]):
            return "你好！我是AI助手，很高兴为您服务。"
        elif any(word in user_message_lower for word in ["天气", "weather"]):
            return "我是一个AI助手，无法获取实时天气信息。建议您查看天气预报应用。"
        elif any(word in user_message_lower for word in ["设备", "device"]):
            return "关于设备管理的问题，我可以帮您解答基本操作和状态查询。"
        elif any(word in user_message_lower for word in ["帮助", "help"]):
            return """我可以帮助您处理以下事务：
1. 设备状态查询
2. 数据报告分析
3. 一般问题解答
4. 技术支持

请告诉我您需要什么帮助？"""
        else:
            return f"我已经收到您的消息：\"{user_message}\"。这是一个模拟的AI回复，在实际应用中会连接真实的AI服务。"


# 全局管理器
manager = ConnectionManager()
ai_service = AIChatService()


@app.websocket("/ws/{device_id}")
async def websocket_endpoint(websocket: WebSocket, device_id: str):
    await manager.connect_device(device_id, websocket)

    try:
        while True:
            # 接收消息
            data = await websocket.receive_text()
            message_dict = json.loads(data)

            # 解析消息类型
            message_type = message_dict.get("type")
            message_data = message_dict.get("data", {})

            # 处理时间戳
            timestamp_str = message_dict.get('timestamp')
            timestamp = datetime.fromisoformat(timestamp_str) if timestamp_str else datetime.now()

            logger.info(f"收到来自设备 {device_id} 的消息: {message_type}")

            if message_type == MessageType.DEVICE_REGISTER.value:
                # 设备注册处理
                device_status = DeviceStatus(
                    deviceId=device_id,
                    versionName=message_data.get("versionName"),
                    url=message_data.get("url"),
                    code=message_data.get("code")
                )

                logger.info(f"设备注册: {device_status}")

                # 发送注册成功响应
                response = {
                    "type": MessageType.STATUS_UPDATE.value,
                    "data": {
                        "status": "registered",
                        "message": "设备注册成功",
                        "deviceId": device_id
                    },
                    "timestamp": datetime.now().isoformat()
                }
                await websocket.send_text(json.dumps(response))

            elif message_type == MessageType.DATA_REPORT.value:
                # 数据报告处理
                content = message_data.get('content', '')
                data_points = message_data.get('data_points', [])

                logger.info(f"设备 {device_id} 数据：{content}")
                logger.info(f"设备 {device_id} 数据报告: {len(data_points)} 个数据点")

                # 处理数据点
                await process_data_points(device_id, data_points)

                # 发送确认响应
                response = {
                    "type": MessageType.STATUS_UPDATE.value,
                    "data": {
                        "status": "data_received",
                        "message": f"成功接收 {len(data_points)} 个数据点",
                        "deviceId": device_id
                    },
                    "timestamp": datetime.now().isoformat()
                }
                await websocket.send_text(json.dumps(response))

            elif message_type == MessageType.AI_CHAT.value:
                # AI对话处理
                user_message = message_data.get('message', '')
                logger.info(f"设备 {device_id} AI对话请求: {user_message}")

                # 处理AI对话
                ai_response = await ai_service.process_chat_message(device_id, user_message, manager)

                # 发送AI回复
                response = {
                    "type": MessageType.AI_CHAT.value,
                    "data": {
                        "role": "assistant",
                        "message": ai_response,
                        "timestamp": datetime.now().isoformat()
                    },
                    "timestamp": datetime.now().isoformat()
                }
                await websocket.send_text(json.dumps(response))

            elif message_type == MessageType.HEARTBEAT.value:
                # 心跳处理
                logger.debug(f"设备 {device_id} 心跳检测")
                response = {
                    "type": MessageType.STATUS_UPDATE.value,
                    "data": {"status": "alive", "deviceId": device_id},
                    "timestamp": datetime.now().isoformat()
                }
                await websocket.send_text(json.dumps(response))

            else:
                logger.warning(f"未知消息类型: {message_type}")
                response = {
                    "type": MessageType.STATUS_UPDATE.value,
                    "data": {"status": "error", "message": "未知的消息类型"},
                    "timestamp": datetime.now().isoformat()
                }
                await websocket.send_text(json.dumps(response))

    except WebSocketDisconnect:
        manager.disconnect_device(device_id)
    except Exception as e:
        logger.error(f"WebSocket处理错误: {str(e)}")
        manager.disconnect_device(device_id)


async def process_data_points(device_id: str, data_points: List):
    """处理数据点"""
    if not data_points:
        return

    logger.info(f"处理设备 {device_id} 的数据点:")
    for i, point in enumerate(data_points):
        logger.info(f"数据点 {i + 1}: {point}")

        # 这里可以添加具体的数据处理逻辑
        # 例如：存储到数据库、进行分析、触发告警等


@app.get("/")
async def root():
    return {"message": "设备管理与AI对话服务"}


@app.get("/status")
async def get_status():
    """获取服务状态"""
    return {
        "active_devices": len(manager.active_device_connections),
        "active_sessions": len(manager.chat_sessions),
        "device_ids": list(manager.active_device_connections.keys())
    }


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)